{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from agents import Agent, WebSearchTool, trace, Runner, gen_trace_id, function_tool, OpenAIChatCompletionsModel\n",
    "from agents.model_settings import ModelSettings\n",
    "from pydantic import BaseModel, Field\n",
    "from dotenv import load_dotenv\n",
    "import asyncio\n",
    "import sendgrid\n",
    "import os\n",
    "from sendgrid.helpers.mail import Mail, Email, To, Content\n",
    "from typing import Dict\n",
    "from IPython.display import display, Markdown\n",
    "from openai import AsyncOpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add ollama model\n",
    "OLLAMA_BASE_URL = \"http://localhost:11434/v1\"\n",
    "ollama_client = AsyncOpenAI(base_url=OLLAMA_BASE_URL, api_key=\"ollama\")\n",
    "ollama_model = OpenAIChatCompletionsModel(model=\"llama3.2\", openai_client=ollama_client)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "INSTRUCTIONS = \"You are a research assistant. Given a search term, you search the web for that term and \\\n",
    "produce a concise summary of the results. The summary must 2-3 paragraphs and less than 300 \\\n",
    "words. Capture the main points. Write succintly, no need to have complete sentences or good \\\n",
    "grammar. This will be consumed by someone synthesizing a report, so it's vital you capture the \\\n",
    "essence and ignore any fluff. Do not include any additional commentary other than the summary itself.\"\n",
    "\n",
    "@function_tool\n",
    "def custom_web_search():\n",
    "    \"\"\" Use the WebSearchTool to search the web for the query \"\"\"\n",
    "    return WebSearchTool(search_context_size=\"low\")\n",
    "\n",
    "search_agent = Agent(\n",
    "    name=\"Search agent\",\n",
    "    instructions=INSTRUCTIONS,\n",
    "    tools=[custom_web_search],\n",
    "    model=ollama_model,\n",
    "    model_settings=ModelSettings(tool_choice=\"required\"),\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "message = \"Latest AI Agent frameworks in 2025\"\n",
    "\n",
    "with trace(\"Search\"):\n",
    "    result = await Runner.run(search_agent, message)\n",
    "\n",
    "display(Markdown(result.final_output))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We will now use Structured Outputs, and include a description of the fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# See note above about cost of WebSearchTool\n",
    "\n",
    "HOW_MANY_SEARCHES = 3\n",
    "\n",
    "INSTRUCTIONS = f\"You are a helpful research assistant. Given a query, come up with a set of web searches \\\n",
    "to perform to best answer the query. Output {HOW_MANY_SEARCHES} terms to query for.\"\n",
    "\n",
    "# Use Pydantic to define the Schema of our response - this is known as \"Structured Outputs\"\n",
    "# With massive thanks to student Wes C. for discovering and fixing a nasty bug with this!\n",
    "\n",
    "class WebSearchItem(BaseModel):\n",
    "    reason: str = Field(description=\"Your reasoning for why this search is important to the query.\")\n",
    "\n",
    "    query: str = Field(description=\"The search term to use for the web search.\")\n",
    "\n",
    "\n",
    "class WebSearchPlan(BaseModel):\n",
    "    searches: list[WebSearchItem] = Field(description=\"A list of web searches to perform to best answer the query.\")\n",
    "\n",
    "\n",
    "planner_agent = Agent(\n",
    "    name=\"PlannerAgent\",\n",
    "    instructions=INSTRUCTIONS,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=WebSearchPlan,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "message = \"Latest AI Agent frameworks in 2025\"\n",
    "\n",
    "with trace(\"Search\"):\n",
    "    result = await Runner.run(planner_agent, message)\n",
    "    print(result.final_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@function_tool\n",
    "def send_email(subject: str, html_body: str) -> Dict[str, str]:\n",
    "    \"\"\" Send out an email with the given subject and HTML body \"\"\"\n",
    "    sg = sendgrid.SendGridAPIClient(api_key=os.environ.get('SENDGRID_API_KEY'))\n",
    "    from_email = Email(\"ed@edwarddonner.com\") # Change this to your verified email\n",
    "    to_email = To(\"ed.donner@gmail.com\") # Change this to your email\n",
    "    content = Content(\"text/html\", html_body)\n",
    "    mail = Mail(from_email, to_email, subject, content).get()\n",
    "    response = sg.client.mail.send.post(request_body=mail)\n",
    "    return {\"status\": \"success\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "send_email"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "INSTRUCTIONS = \"\"\"You are able to send a nicely formatted HTML email based on a detailed report.\n",
    "You will be provided with a detailed report. You should use your tool to send one email, providing the \n",
    "report converted into clean, well presented HTML with an appropriate subject line.\"\"\"\n",
    "\n",
    "email_agent = Agent(\n",
    "    name=\"Email agent\",\n",
    "    instructions=INSTRUCTIONS,\n",
    "    tools=[send_email],\n",
    "    model=\"gpt-4o-mini\",\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "INSTRUCTIONS = (\n",
    "    \"You are a senior researcher tasked with writing a cohesive report for a research query. \"\n",
    "    \"You will be provided with the original query, and some initial research done by a research assistant.\\n\"\n",
    "    \"You should first come up with an outline for the report that describes the structure and \"\n",
    "    \"flow of the report. Then, generate the report and return that as your final output.\\n\"\n",
    "    \"The final output should be in markdown format, and it should be lengthy and detailed. Aim \"\n",
    "    \"for 5-10 pages of content, at least 1000 words.\"\n",
    ")\n",
    "\n",
    "\n",
    "class ReportData(BaseModel):\n",
    "    short_summary: str = Field(description=\"A short 2-3 sentence summary of the findings.\")\n",
    "\n",
    "    markdown_report: str = Field(description=\"The final report\")\n",
    "\n",
    "    follow_up_questions: list[str] = Field(description=\"Suggested topics to research further\")\n",
    "\n",
    "\n",
    "writer_agent = Agent(\n",
    "    name=\"WriterAgent\",\n",
    "    instructions=INSTRUCTIONS,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=ReportData,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The next 3 functions will plan and execute the search, using planner_agent and search_agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def plan_searches(query: str):\n",
    "    \"\"\" Use the planner_agent to plan which searches to run for the query \"\"\"\n",
    "    print(\"Planning searches...\")\n",
    "    result = await Runner.run(planner_agent, f\"Query: {query}\")\n",
    "    print(f\"Will perform {len(result.final_output.searches)} searches\")\n",
    "    return result.final_output\n",
    "\n",
    "async def perform_searches(search_plan: WebSearchPlan):\n",
    "    \"\"\" Call search() for each item in the search plan \"\"\"\n",
    "    print(\"Searching...\")\n",
    "    tasks = [asyncio.create_task(search(item)) for item in search_plan.searches]\n",
    "    results = await asyncio.gather(*tasks)\n",
    "    print(\"Finished searching\")\n",
    "    return results\n",
    "\n",
    "async def search(item: WebSearchItem):\n",
    "    \"\"\" Use the search agent to run a web search for each item in the search plan \"\"\"\n",
    "    input = f\"Search term: {item.query}\\nReason for searching: {item.reason}\"\n",
    "    result = await Runner.run(search_agent, input)\n",
    "    return result.final_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The next 2 functions write a report and email it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def write_report(query: str, search_results: list[str]):\n",
    "    \"\"\" Use the writer agent to write a report based on the search results\"\"\"\n",
    "    print(\"Thinking about report...\")\n",
    "    input = f\"Original query: {query}\\nSummarized search results: {search_results}\"\n",
    "    result = await Runner.run(writer_agent, input)\n",
    "    print(\"Finished writing report\")\n",
    "    return result.final_output\n",
    "\n",
    "async def send_email(report: ReportData):\n",
    "    \"\"\" Use the email agent to send an email with the report \"\"\"\n",
    "    print(\"Writing email...\")\n",
    "    result = await Runner.run(email_agent, report.markdown_report)\n",
    "    print(\"Email sent\")\n",
    "    return report"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Showtime!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query =\"Latest AI Agent frameworks in 2025\"\n",
    "\n",
    "with trace(\"Research trace\"):\n",
    "    print(\"Starting research...\")\n",
    "    search_plan = await plan_searches(query)\n",
    "    search_results = await perform_searches(search_plan)\n",
    "    report = await write_report(query, search_results)\n",
    "    await send_email(report)  \n",
    "    print(\"Hooray!\")\n",
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
